// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <nebula/csv/column_decoder.h>

#include <cstddef>
#include <cstdint>
#include <memory>
#include <sstream>
#include <string>
#include <utility>

#include <nebula/core/array.h>
#include <nebula/array/builder_base.h>
#include <nebula/csv/converter.h>
#include <nebula/csv/inference_internal.h>
#include <nebula/csv/options.h>
#include <nebula/csv/parser.h>
#include <turbo/utility/status.h>
#include <nebula/types/type.h>
#include <nebula/types/type_fwd.h>
#include <nebula/future/future.h>
#include <turbo/log/logging.h>

namespace nebula::csv {

    class ConcreteColumnDecoder : public ColumnDecoder {
    public:
        explicit ConcreteColumnDecoder(MemoryPool *pool, int32_t col_index = -1)
                : ColumnDecoder(), pool_(pool), col_index_(col_index) {}

    protected:
        // XXX useful?
        virtual std::shared_ptr<DataType> type() const = 0;

        turbo::Result<std::shared_ptr<Array>> WrapConversionError(
                const turbo::Result<std::shared_ptr<Array>> &result) {
            if (TURBO_LIKELY(result.ok())) {
                return result;
            } else {
                const auto &st = result.status();
                std::stringstream ss;
                ss << "In CSV column #" << col_index_ << ": " << st.message();
                return st.with_message(ss.str());
            }
        }

        MemoryPool *pool_;
        int32_t col_index_;
    };

//////////////////////////////////////////////////////////////////////////
// Null column decoder implementation (for a column not in the CSV file)

    class NullColumnDecoder : public ConcreteColumnDecoder {
    public:
        explicit NullColumnDecoder(const std::shared_ptr<DataType> &type, MemoryPool *pool)
                : ConcreteColumnDecoder(pool), type_(type) {}

        Future<std::shared_ptr<Array>> Decode(
                const std::shared_ptr<BlockParser> &parser) override;

    protected:
        std::shared_ptr<DataType> type() const override { return type_; }

        std::shared_ptr<DataType> type_;
    };

    Future<std::shared_ptr<Array>> NullColumnDecoder::Decode(
            const std::shared_ptr<BlockParser> &parser) {
                DKCHECK_GE(parser->num_rows(), 0);
        return WrapConversionError(MakeArrayOfNull(type_, parser->num_rows(), pool_));
    }

//////////////////////////////////////////////////////////////////////////
// Pre-typed column decoder implementation

    class TypedColumnDecoder : public ConcreteColumnDecoder {
    public:
        TypedColumnDecoder(const std::shared_ptr<DataType> &type, int32_t col_index,
                           const ConvertOptions &options, MemoryPool *pool)
                : ConcreteColumnDecoder(pool, col_index), type_(type), options_(options) {}

        turbo::Status init();

        Future<std::shared_ptr<Array>> Decode(
                const std::shared_ptr<BlockParser> &parser) override;

    protected:
        std::shared_ptr<DataType> type() const override { return type_; }

        std::shared_ptr<DataType> type_;
        // CAUTION: ConvertOptions can grow large (if it customizes hundreds or
        // thousands of columns), so avoid copying it in each TypedColumnDecoder.
        const ConvertOptions &options_;

        std::shared_ptr<Converter> converter_;
    };

    turbo::Status TypedColumnDecoder::init() {
        TURBO_MOVE_OR_RAISE(converter_, Converter::create(type_, options_, pool_));
        return turbo::OkStatus();
    }

    Future<std::shared_ptr<Array>> TypedColumnDecoder::Decode(
            const std::shared_ptr<BlockParser> &parser) {
                DKCHECK_NE(converter_, nullptr);
        return Future<std::shared_ptr<Array>>::make_finished(
                WrapConversionError(converter_->Convert(*parser, col_index_)));
    }

//////////////////////////////////////////////////////////////////////////
// Type-inferring column builder implementation

    class InferringColumnDecoder : public ConcreteColumnDecoder {
    public:
        InferringColumnDecoder(int32_t col_index, const ConvertOptions &options,
                               MemoryPool *pool)
                : ConcreteColumnDecoder(pool, col_index),
                  options_(options),
                  infer_status_(options),
                  type_frozen_(false) {
            first_inference_run_ = Future<>::create();
            first_inferrer_ = 0;
        }

        turbo::Status init();

        Future<std::shared_ptr<Array>> Decode(
                const std::shared_ptr<BlockParser> &parser) override;

    protected:
        std::shared_ptr<DataType> type() const override {
                    DKCHECK_NE(converter_, nullptr);
            return converter_->type();
        }

        turbo::Status UpdateType();

        turbo::Result<std::shared_ptr<Array>> RunInference(const std::shared_ptr<BlockParser> &parser);

        // CAUTION: ConvertOptions can grow large (if it customizes hundreds or
        // thousands of columns), so avoid copying it in each InferringColumnDecoder.
        const ConvertOptions &options_;

        // Current inference status
        InferStatus infer_status_;
        bool type_frozen_;
        std::atomic<int> first_inferrer_;
        Future<> first_inference_run_;
        std::shared_ptr<Converter> converter_;
    };

    turbo::Status InferringColumnDecoder::init() { return UpdateType(); }

    turbo::Status InferringColumnDecoder::UpdateType() {
        return infer_status_.make_converter(pool_).try_value(&converter_);
    }

    turbo::Result<std::shared_ptr<Array>> InferringColumnDecoder::RunInference(
            const std::shared_ptr<BlockParser> &parser) {
        while (true) {
            // (no one else should be updating converter_ concurrently)
            auto maybe_array = converter_->Convert(*parser, col_index_);

            if (maybe_array.ok() || !infer_status_.can_loosen_type()) {
                // Conversion succeeded, or failed definitively
                        DKCHECK(!type_frozen_);
                type_frozen_ = true;
                return maybe_array;
            }
            // Conversion failed temporarily, try another type
            infer_status_.LoosenType(maybe_array.status());
            auto update_status = UpdateType();
            if (!update_status.ok()) {
                return update_status;
            }
        }
    }

    Future<std::shared_ptr<Array>> InferringColumnDecoder::Decode(
            const std::shared_ptr<BlockParser> &parser) {
        // Empty arrays before the first inference run must be discarded since the type of the
        // array will be NA and not match arrays decoded later
        if (parser->num_rows() == 0) {
            return Future<std::shared_ptr<Array>>::make_finished(
                    MakeArrayOfNull(converter_->type(), 0));
        }

        bool already_taken = first_inferrer_.fetch_or(1);
        // First block: run inference
        if (!already_taken) {
            auto maybe_array = RunInference(parser);
            first_inference_run_.mark_finished();
            return Future<std::shared_ptr<Array>>::make_finished(std::move(maybe_array));
        }

        // Non-first block: wait for inference to finish on first block now,
        // without blocking a worker thread.
        return first_inference_run_.Then([this, parser] {
                    DKCHECK(type_frozen_);
            auto maybe_array = converter_->Convert(*parser, col_index_);
            return WrapConversionError(converter_->Convert(*parser, col_index_));
        });
    }

    //////////////////////////////////////////////////////////////////////////
    // Factory functions

    turbo::Result<std::shared_ptr<ColumnDecoder>> ColumnDecoder::create(
            MemoryPool *pool, int32_t col_index, const ConvertOptions &options) {
        auto ptr = std::make_shared<InferringColumnDecoder>(col_index, options, pool);
        TURBO_RETURN_NOT_OK(ptr->init());
        return ptr;
    }

    turbo::Result<std::shared_ptr<ColumnDecoder>> ColumnDecoder::create(
            MemoryPool *pool, std::shared_ptr<DataType> type, int32_t col_index,
            const ConvertOptions &options) {
        auto ptr =
                std::make_shared<TypedColumnDecoder>(std::move(type), col_index, options, pool);
        TURBO_RETURN_NOT_OK(ptr->init());
        return ptr;
    }

    turbo::Result<std::shared_ptr<ColumnDecoder>> ColumnDecoder::MakeNull(
            MemoryPool *pool, std::shared_ptr<DataType> type) {
        return std::make_shared<NullColumnDecoder>(std::move(type), pool);
    }

}  // namespace nebula::csv
